"""TensorRT引擎构建管理器"""
from typing import List, Optional
from core.tensorrt_core import TensorRTBuilder, QuantizationStrategy
from core.ppq_quantizer import PPQINT8Strategy

class EngineBuilder(TensorRTBuilder):
    """TensorRT引擎构建管理器"""
    
    def __init__(self, workspace_size: int = 1 << 32):
        super().__init__(workspace_size)
        self.quantization_strategies: List[QuantizationStrategy] = []
    
    def add_quantization_strategy(self, strategy: QuantizationStrategy) -> 'EngineBuilder':
        """添加量化策略"""
        self.quantization_strategies.append(strategy)
        return self
    
    def build_engine(self, onnx_path: str, engine_path: str) -> None:
        """构建TensorRT引擎"""
        # 创建网络和配置
        self._create_network()
        self._create_config()
        
        # 解析ONNX
        self._parse_onnx(onnx_path)
        
        # 应用量化策略
        for strategy in self.quantization_strategies:
            strategy.apply_quantization(self.config, self.network)
        
        # 构建并保存引擎
        serialized_engine = self._build_and_serialize()
        self.save_engine(serialized_engine, engine_path)

class PPQEngineBuilder:
    """PPQ专用引擎构建器"""
    
    @staticmethod
    def build_int8_engine(
        onnx_path: str,
        ppq_json_path: str,
        engine_path: str,
        workspace_size: int = 1 << 32,
        enable_fp16: bool = False
    ) -> None:
        """构建PPQ INT8引擎"""
        builder = EngineBuilder(workspace_size)
        
        # 添加量化策略
        if enable_fp16:
            from ..core.tensorrt_core import FP16Strategy
            builder.add_quantization_strategy(FP16Strategy())
        
        builder.add_quantization_strategy(PPQINT8Strategy(ppq_json_path))
        
        # 构建引擎
        builder.build_engine(onnx_path, engine_path)